package com.example.registry;

import com.example.config.ShardRabbitListenerConfigUtils;
import com.example.listener.MethodShardRabbitListenerEndpoint;
import com.example.lock.ShardRabbitLock;
import com.example.lock.ShardRabbitLockManager;
import org.springframework.amqp.rabbit.listener.MessageListenerContainer;
import org.springframework.amqp.rabbit.listener.RabbitListenerContainerFactory;
import org.springframework.amqp.rabbit.listener.RabbitListenerEndpoint;
import org.springframework.amqp.rabbit.listener.RabbitListenerEndpointRegistry;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.BeanFactory;
import org.springframework.beans.factory.BeanFactoryAware;
import org.springframework.data.redis.core.StringRedisTemplate;

public class ShardRabbitListenerEndpointRegistry extends RabbitListenerEndpointRegistry implements BeanFactoryAware {
    private BeanFactory beanFactory;

    public ShardRabbitListenerEndpointRegistry() {
    }

    public void setBeanFactory(BeanFactory beanFactory) throws BeansException {
        this.beanFactory = beanFactory;
    }

    public void registerListenerContainer(RabbitListenerEndpoint endpoint, RabbitListenerContainerFactory<?> factory, boolean startImmediately) {
        this.logger.info("force to startImmediately to false");
        super.registerListenerContainer(endpoint, factory, false);
        MessageListenerContainer messageListenerContainer = this.getListenerContainer(endpoint.getId());
        ShardRabbitLockManager shardRabbitLockManager = (ShardRabbitLockManager)this.beanFactory.getBean(ShardRabbitListenerConfigUtils.RABBIT_LISTENER_LOCK_MANAGER_BEAN_NAME, ShardRabbitLockManager.class);
        StringRedisTemplate stringRedisTemplate = (StringRedisTemplate)this.beanFactory.getBean("stringRedisTemplate", StringRedisTemplate.class);
        MethodShardRabbitListenerEndpoint methodShardRabbitListenerEndpoint = (MethodShardRabbitListenerEndpoint)endpoint;
        ShardRabbitLock shardRabbitLock = new ShardRabbitLock();
        shardRabbitLock.setQueueNamePrefix(methodShardRabbitListenerEndpoint.getQueueNamePrefix());
        shardRabbitLock.setSize(methodShardRabbitListenerEndpoint.getShardSize());
        shardRabbitLock.setNum(methodShardRabbitListenerEndpoint.getNum());
        shardRabbitLock.setContainer(messageListenerContainer);
        shardRabbitLock.setRedisTemplate(stringRedisTemplate);
        shardRabbitLock.setInstanceId(methodShardRabbitListenerEndpoint.getInstanceId());
        shardRabbitLockManager.add(shardRabbitLock);
    }
}
